// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>

#include <pollux/testing/gtest_utils.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux::test {

class MockVectorSerde : public VectorSerde {
 public:
  MockVectorSerde() : VectorSerde(VectorSerde::Kind::kPresto) {}

  void estimateSerializedSize(
      const BaseVector* /*vector*/,
      const melon::Range<const IndexRange*>& ranges,
      vector_size_t** sizes) override {}

  std::unique_ptr<IterativeVectorSerializer> createIterativeSerializer(
      RowTypePtr type,
      int32_t numRows,
      StreamArena* streamArena,
      const Options* options = nullptr) override {
    return nullptr;
  };

  void deserialize(
      ByteInputStream* source,
      pollux::memory::MemoryPool* pool,
      RowTypePtr type,
      RowVectorPtr* result,
      const Options* options = nullptr) override {}
};

TEST(VectorStreamTest, serdeRegistration) {
  deregisterVectorSerde();

  // Nothing registered yet.
  EXPECT_FALSE(isRegisteredVectorSerde());
  EXPECT_THROW(getVectorSerde(), PolluxRuntimeError);

  // Register a mock serde.
  registerVectorSerde(std::make_unique<MockVectorSerde>());

  EXPECT_TRUE(isRegisteredVectorSerde());
  auto serde = getVectorSerde();
  EXPECT_NE(serde, nullptr);
  EXPECT_NE(dynamic_cast<MockVectorSerde*>(serde), nullptr);

  // Can't double register.
  EXPECT_THROW(
      registerVectorSerde(std::make_unique<MockVectorSerde>()),
      PolluxRuntimeError);

  deregisterVectorSerde();
  EXPECT_FALSE(isRegisteredVectorSerde());
}

TEST(VectorStreamTest, namedSerdeRegistration) {
  const VectorSerde::Kind kind = VectorSerde::Kind::kPresto;

  // Nothing registered yet.
  deregisterNamedVectorSerde(kind);
  EXPECT_FALSE(isRegisteredNamedVectorSerde(kind));
  POLLUX_ASSERT_THROW(
      getNamedVectorSerde(kind),
      "Named vector serde 'Presto' is not registered.");

  // Register a mock serde.
  registerNamedVectorSerde(kind, std::make_unique<MockVectorSerde>());

  auto serde = getNamedVectorSerde(kind);
  EXPECT_NE(serde, nullptr);
  EXPECT_NE(dynamic_cast<MockVectorSerde*>(serde), nullptr);

  const VectorSerde::Kind otherKind = VectorSerde::Kind::kUnsafeRow;
  EXPECT_FALSE(isRegisteredNamedVectorSerde(otherKind));
  POLLUX_ASSERT_THROW(
      getNamedVectorSerde(otherKind),
      "Named vector serde 'UnsafeRow' is not registered.");

  // Can't double register.
  POLLUX_ASSERT_THROW(
      registerNamedVectorSerde(kind, std::make_unique<MockVectorSerde>()),
      "Vector serde 'Presto' is already registered.");

  // Register another one.
  EXPECT_FALSE(isRegisteredNamedVectorSerde(otherKind));
  EXPECT_THROW(getNamedVectorSerde(otherKind), PolluxRuntimeError);
  registerNamedVectorSerde(otherKind, std::make_unique<MockVectorSerde>());
  EXPECT_TRUE(isRegisteredNamedVectorSerde(otherKind));

  deregisterNamedVectorSerde(otherKind);
  EXPECT_FALSE(isRegisteredNamedVectorSerde(otherKind));
}
} // namespace kumo::pollux::test
